from alive_progress import alive_bar
from dataclasses import dataclass
import time
from typing import TYPE_CHECKING
from multiprocessing import Pool, Manager
from rich import print

from solver.solver import Solution, solve
from util.load_puzzles import load_from_file
from util.SConvert import make_blank, make_sboard, make_scomponents
from container import *

if TYPE_CHECKING:
    from container import nonBlankLabel_t
    from puzzle.board import Board
    from solver.SComponent import SComponent
    from multiprocessing.synchronize import Lock


def _solve(
    board: "Board",
    sc: "list[list[SComponent]]",
    month: "nonBlankLabel_t",
    day: "nonBlankLabel_t",
    weekday: "nonBlankLabel_t",
    solve_all: bool = False
) -> list[Solution]:
    sb = make_sboard(board)
    
    make_blank(sb, board.get_td_by_label(month))
    make_blank(sb, board.get_td_by_label(day))
    make_blank(sb, board.get_td_by_label(weekday))
    
    solutions = solve(sb, sc, solve_all)
    
    return solutions


@dataclass
class TestResult:
    month: "nonBlankLabel_t"
    day: "nonBlankLabel_t"
    weekday: "nonBlankLabel_t"
    n_solutions: int
    first_it: int
    tot_it: int
    time_ellapsed: float

def _test_all_task(
    lock: "Lock",
    board: "Board",
    sc: "list[list[SComponent]]",
    month: "nonBlankLabel_t",
    day: "nonBlankLabel_t",
    weekday: "nonBlankLabel_t",
    output: bool,
) -> TestResult:
    st = time.time()
    solutions = _solve(board, sc, month, day, weekday, True)
    et = time.time()
    
    tr = TestResult(
        month, day, weekday,
        len(solutions),
        solutions[0].it,
        sum([sol.it for sol in solutions]),
        round(et - st, 2)
    )
    
    if output:
        with lock:
            print("|{0:^5}|{1:^4}|{2:^5}| {3:<6}| {4:<5}| {5:<11}".format(
                month, day, weekday,
                str(tr.n_solutions),
                str(tr.first_it), str(tr.tot_it)
            ))

    return tr

def _validate_task(
    lock: "Lock",
    board: "Board",
    sc: "list[list[SComponent]]",
    month: "nonBlankLabel_t",
    day: "nonBlankLabel_t",
    weekday: "nonBlankLabel_t"
) -> None:
    solutions = _solve(board, sc, month, day, weekday, False)
    if len(solutions) == 0:
        with lock:
            print(f"Failed {month} {day} {weekday}")

if __name__ == "__main__":
    
    board, _ = load_from_file("Final Design", "./data/des_final.txt", validation = False)
    _, comps = load_from_file("", "./data/des_ref.txt")
    
    print(f"Working on board {board.name}")
    
    sc = []
    for transforms in comps:
        sc.append(make_scomponents(transforms))
    
    if True:
        
        lock = Manager().Lock()
        p = Pool(8)
        
        for m in LABEL_MONTH:
            for d in LABEL_DAY:
                for w in LABEL_WEEKDAY:
                    if m in ["Feb", "Apr", "Jun", "Sep", "Nov"] and d == "31":
                        continue
                    if m == "Feb" and d == "30":
                        continue
                    
                    p.apply_async(
                        _validate_task,
                        args = (lock, board, sc, m, d, w, ),
                    )
            
        p.close()
        p.join()
        
        print("Validation done")
    
    if False:
        
        lock = Manager().Lock()
        p = Pool(8)
        
        for m in LABEL_MONTH:
            for d in LABEL_DAY:
                for w in LABEL_WEEKDAY:
                    if m in ["Feb", "Apr", "Jun", "Sep", "Nov"] and d == "31":
                        continue
                    if m == "Feb" and d == "30":
                        continue
                    
                    p.apply_async(
                        _test_all_task,
                        args = (lock, board, sc, m, d, w, True, ),
                    )
            
        p.close()
        p.join()